%%--- coding:utf-8 ---
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%% File Name: mc_auth_logic
%%% Created on : 2024/4/19 8:55
%%% @author Gaylen 252323463@qq.com
%%% @copyright (C) 2024, freedom
%%% @doc
%%%
%%% @end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

-module(mc_auth_logic).
-author("Gaylen").
-include("mongodb_driver.hrl").

%% API
-export([
    auth/4
]).

-define(RANDOM_LENGTH, 24).
-define(AUTH_CMD(User, Nonce, Password),
    {
        <<"authenticate">>, 1,
        <<"user">>, User,
        <<"nonce">>, Nonce,
        <<"key">>, mc_utils:pw_key(Nonce, User, Password)
    }).

%% Authorize on database synchronously
-spec auth(pid(), database(), binary() | undefined, binary() | undefined) -> boolean().
auth(Connection, Database, User, Password) ->
    priv_scram_sha_1_auth(Connection, Database, User, Password).

%% @private
-spec priv_scram_sha_1_auth(port(), binary(), binary(), binary()) -> boolean().
priv_scram_sha_1_auth(Connection, Database, User, Password) ->
    priv_scram_first_step(Connection, Database, User, Password).

%% @private
priv_scram_first_step(Connection, Database, User, Password) ->
    RandomBString = mc_utils:random_nonce(?RANDOM_LENGTH),
    FirstMessage = priv_compose_first_message(User, RandomBString),
    Message = <<?GS2_HEADER/binary, FirstMessage/binary>>,
    Cmd = {<<"saslStart">>, 1, <<"mechanism">>, <<"SCRAM-SHA-1">>, <<"payload">>, {bin, bin, Message}, <<"autoAuthorize">>, 1},
    {true, Res} = mc_worker_api_logic:database_command(Connection, Database, Cmd),
    ConversationId = maps:get(<<"conversationId">>, Res, {}),
    Payload = maps:get(<<"payload">>, Res),
    priv_scram_second_step(Connection, Database, User, Password, Payload, ConversationId, RandomBString, FirstMessage).

%% @private
priv_compose_first_message(User, RandomBString) ->
    UserName = <<<<"n=">>/binary, (mc_utils:encode_name(User))/binary>>,
    Nonce = <<<<"r=">>/binary, RandomBString/binary>>,
    <<UserName/binary, <<",">>/binary, Nonce/binary>>.

%% @private
priv_scram_second_step(Connection, Database, Login, Password, {bin, bin, Decoded} = _Payload, ConversationId, RandomBString, FirstMessage) ->
    {Signature, ClientFinalMessage} = priv_compose_second_message(Decoded, Login, Password, RandomBString, FirstMessage),
    Cmd = {<<"saslContinue">>, 1, <<"conversationId">>, ConversationId,
        <<"payload">>, {bin, bin, ClientFinalMessage}},
    {true, Res} = mc_worker_api_logic:database_command(Connection, Database, Cmd),
    priv_scram_third_step(Connection, base64:encode(Signature), Res, ConversationId, Database).

%% @private
priv_compose_second_message(Payload, Login, Password, RandomBString, FirstMessage) ->
    ParamList = priv_parse_server_responce(Payload),
    R = mc_utils:get_value(<<"r">>, ParamList),
    Nonce = <<<<"r=">>/binary, R/binary>>,
    {0, ?RANDOM_LENGTH} = binary:match(R, [RandomBString], []),
    S = mc_utils:get_value(<<"s">>, ParamList),
    I = binary_to_integer(mc_utils:get_value(<<"i">>, ParamList)),
    SaltedPassword = priv_salt_pwd(mc_utils:pw_hash(Login, Password), base64:decode(S), I),
    ChannelBinding = <<<<"c=">>/binary, (base64:encode(?GS2_HEADER))/binary>>,
    ClientFinalMessageWithoutProof = <<ChannelBinding/binary, <<",">>/binary, Nonce/binary>>,
    AuthMessage = <<FirstMessage/binary, <<",">>/binary, Payload/binary, <<",">>/binary, ClientFinalMessageWithoutProof/binary>>,
    ServerSignature = priv_generate_sig(SaltedPassword, AuthMessage),
    Proof = priv_generate_proof(SaltedPassword, AuthMessage),
    {ServerSignature, <<ClientFinalMessageWithoutProof/binary, <<",">>/binary, Proof/binary>>}.

%% @private
priv_parse_server_responce(Response) ->
    ParamList = binary:split(Response, <<",">>, [global]),
    lists:map(
        fun(Param) ->
            [K, V] = binary:split(Param, <<"=">>),
            {K, V}
        end, ParamList).

%% @private
priv_salt_pwd(Password, Salt, Iterations) ->
    crypto:pbkdf2_hmac(sha, Password, Salt, Iterations, 20).

%%priv_pbkdf2(MacFunc, Password, Salt, Iterations, DerivedLength) ->
%%    MacFunc1 = priv_resolve_mac_func(MacFunc),
%%    Bin = priv_pbkdf2_loop(MacFunc1, Password, Salt, Iterations, DerivedLength, 1, []),
%%    {ok, Bin}.
%%
%%priv_pbkdf2_loop(MacFunc, Password, Salt, Iterations, DerivedLength, BlockIndex, Acc) ->
%%    case iolist_size(Acc) > DerivedLength of
%%        true ->
%%            <<Bin:DerivedLength/binary, _/binary>> = iolist_to_binary(lists:reverse(Acc)),
%%            Bin;
%%        false ->
%%            Block = priv_pbkdf2_loop2(MacFunc, Password, Salt, Iterations, BlockIndex, 1, <<>>, <<>>),
%%            priv_pbkdf2_loop(MacFunc, Password, Salt, Iterations, DerivedLength, BlockIndex + 1, [Block | Acc])
%%    end.
%%
%%priv_pbkdf2_loop2(_MacFunc, _Password, _Salt, Iterations, _BlockIndex, Iteration, _Prev, Acc) when Iteration > Iterations ->
%%    Acc;
%%priv_pbkdf2_loop2(MacFunc, Password, Salt, Iterations, BlockIndex, 1, _Prev, _Acc) ->
%%    InitialBlock = MacFunc(Password, <<Salt/binary, BlockIndex:32/integer>>),
%%    priv_pbkdf2_loop2(MacFunc, Password, Salt, Iterations, BlockIndex, 2, InitialBlock, InitialBlock);
%%priv_pbkdf2_loop2(MacFunc, Password, Salt, Iterations, BlockIndex, Iteration, Prev, Acc) ->
%%    Next = MacFunc(Password, Prev),
%%    priv_pbkdf2_loop2(MacFunc, Password, Salt, Iterations, BlockIndex, Iteration + 1, Next, crypto:exor(Next, Acc)).
%%
%%priv_resolve_mac_func({hmac, DigestFunc}) ->
%%    fun(Key, Data) ->
%%        %crypto:hmac(DigestFunc, Key, Data)
%%        HMAC = crypto:pbkdf2_hmac()
%%        HMAC = crypto:hmac_init(DigestFunc, Key),
%%        HMAC1 = crypto:hmac_update(HMAC, Data),
%%        crypto:hmac_final(HMAC1)
%%    end;
%%priv_resolve_mac_func(MacFunc) when is_function(MacFunc) ->
%%    MacFunc;
%%priv_resolve_mac_func(md4) -> priv_resolve_mac_func({hmac, md4});
%%priv_resolve_mac_func(md5) -> priv_resolve_mac_func({hmac, md5});
%%priv_resolve_mac_func(ripemd160) -> priv_resolve_mac_func({hmac, ripemd160});
%%priv_resolve_mac_func(sha) -> priv_resolve_mac_func({hmac, sha});
%%priv_resolve_mac_func(sha224) -> priv_resolve_mac_func({hmac, sha224});
%%priv_resolve_mac_func(sha256) -> priv_resolve_mac_func({hmac, sha256});
%%priv_resolve_mac_func(sha384) -> priv_resolve_mac_func({hmac, sha384});
%%priv_resolve_mac_func(sha512) -> priv_resolve_mac_func({hmac, sha512}).

%% @private
priv_generate_sig(SaltedPassword, AuthMessage) ->
    ServerKey = mc_utils:hmac(SaltedPassword, "Server Key"),
    mc_utils:hmac(ServerKey, AuthMessage).

%% @private
priv_generate_proof(SaltedPassword, AuthMessage) ->
    ClientKey = mc_utils:hmac(SaltedPassword, <<"Client Key">>),
    StoredKey = crypto:hash(sha, ClientKey),
    Signature = mc_utils:hmac(StoredKey, AuthMessage),
    ClientProof = priv_xorKeys(ClientKey, Signature, <<>>),
    <<<<"p=">>/binary, (base64:encode(ClientProof))/binary>>.

%% @private
priv_xorKeys(<<>>, _, Res) -> Res;
priv_xorKeys(<<FA, RestA/binary>>, <<FB, RestB/binary>>, Res) ->
    priv_xorKeys(RestA, RestB, <<Res/binary, <<(FA bxor FB)>>/binary>>).

%% @private
priv_scram_third_step(Connection, ServerSignature, Response, ConversationId, Database) ->
    {bin, bin, Payload} = maps:get(<<"payload">>, Response),
    Done = maps:get(<<"done">>, Response, false),
    ParamList = priv_parse_server_responce(Payload),
    ServerSignature = mc_utils:get_value(<<"v">>, ParamList),
    priv_scram_forth_step(Connection, Done, ConversationId, Database).

%% @private
priv_scram_forth_step(_, true, _, _) -> true;
priv_scram_forth_step(Connection, false, ConversationId, Database) ->
    Cmd = {<<"saslContinue">>, 1, <<"conversationId">>,
        ConversationId, <<"payload">>, {bin, bin, <<>>}},
    {true, Res} = mc_worker_api_logic:database_command(Connection, Database, Cmd),
    true = maps:get(<<"done">>, Res, false).

%%%% @private
%%-spec mongodb_cr_auth(pid(), binary(), binary(), binary()) -> boolean().
%%mongodb_cr_auth(Connection, Database, Login, Password) ->
%%    {true, Res} = mc_worker_api_logic:database_command(Connection, Database, {<<"getnonce">>, 1}),
%%    Nonce = maps:get(<<"nonce">>, Res),
%%    case mc_worker_api_logic:database_command(Connection, Database, ?AUTH_CMD(Login, Nonce, Password)) of
%%        {true, _} -> true;
%%        {false, Reason} -> erlang:error(Reason)
%%    end.